Skip to content

feat(integration): add Pi Agent sandbox integration example (#698)#786

Open
xyaohubery wants to merge 9 commits into
TencentCloud:masterfrom
xyaohubery:feat/pi-agent-integration
Open

feat(integration): add Pi Agent sandbox integration example (#698)#786
xyaohubery wants to merge 9 commits into
TencentCloud:masterfrom
xyaohubery:feat/pi-agent-integration

Conversation

@xyaohubery

Copy link
Copy Markdown

Summary

Add Pi Agent + CubeSandbox integration example following the same 'Sandbox as Tool' pattern as the Claude Code integration.

Contents

  • pi_agent_sandbox_demo.py: Bridge code using cubesandbox SDK to run AI-generated code in hardware-isolated MicroVMs
  • README.md: Architecture overview and quick start instructions

Pattern

Pi Agent runs locally (handling conversation and orchestration) and delegates code execution to CubeSandbox on demand via the Python SDK.

Fixes #698

Assisted-by: claude-code:deepseek-v4-pro

- Add MCP Server exposing 9 CubeSandbox tools via stdio
- Input validation with bounds checking on all numeric/temporal params
- Defensive error handling throughout tool handler
- Bilingual README (EN/ZH) with setup flow and troubleshooting
- Integration guide (EN/ZH) following docs template conventions
- Document WSL2 incompatibility (Cubelet + network-agent) per v0.5.0 testing

Assisted-by: claude-code:deepseek-v4-pro
- Replace silent except:pass in run_code with [mcp: error ...] warnings
- Remove redundant get_info() HTTP call after sandbox creation
- Fix falsy string trap in sandbox_snapshot name handling

Assisted-by: claude-code:deepseek-v4-pro
Critical fixes:
- Use CUBE_API_URL instead of E2B_API_URL (matches SDK Config class)
- Remove non-functional E2B_API_KEY references
- Wrap sync SDK calls via asyncio.to_thread() to avoid blocking event loop
- Add asyncio.Lock to prevent TOCTOU race on _sandbox lifecycle

Medium fixes:
- Sanitize file paths (reject .. traversal)
- Register atexit cleanup to kill sandbox on exit
- Sanitize exception messages (no internal URLs/stack traces)
- Fix sandbox_pause success message (remove non-existent resume flow)
- Remove redundant try/except around Execution dataclass field access
- Align writable layer size to 2G across all docs

Low fixes:
- Move SSL_CERT_FILE side effect from import time to main()
- Change run_code timeout schema from 'integer' to 'number'

Assisted-by: claude-code:deepseek-v4-pro
…loud#698)

- Bridge code for running Pi Agent code in CubeSandbox MicroVMs
- README with architecture and quick start instructions
- Follows same 'Sandbox as Tool' pattern as Claude Code integration

Fixes TencentCloud#698

Assisted-by: claude-code:deepseek-v4-pro
pass
async with _lock:
try:
_sandbox.kill()

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: TOCTOU race — uses _sandbox.kill() instead of sb.kill()

Line 524 calls _sandbox.kill() on the global reference, but the locally captured sb (line 348) is the correct sandbox reference. Between the guard check (lines 347-349, which releases the lock) and this line's lock re-acquisition, a concurrent sandbox_destroy call could have set _sandbox = None (line 526, 529), causing AttributeError. If _sandbox were reassigned by a concurrent sandbox_create, the wrong sandbox would be killed.

Replace _sandbox.kill() with sb.kill().

loop = asyncio.get_running_loop()
for sig in (signal.SIGTERM, signal.SIGINT):
try:
loop.add_signal_handler(sig, lambda: None)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: No-op signal handlers prevent clean shutdown

The installed handler lambda: None silently consumes SIGTERM and SIGINT instead of initiating shutdown. This means:

  1. SIGINT no longer raises KeyboardInterrupt, so asyncio.run() cannot perform graceful cancellation.
  2. SIGTERM no longer terminates the process — the default action would do so.
  3. The atexit-registered _cleanup() never fires, so the sandbox MicroVM remains running as an orphan until its timeout expires (up to 86,400 seconds).

The comment on line 551 claims this forwards signals for "clean shutdown + atexit cleanup", but the no-op achieves the opposite. Replace with a handler that actually stops the event loop, e.g. loop.add_signal_handler(sig, loop.stop), or remove the handlers entirely.

return str(p)


def _sanitize_error(exc: BaseException) -> str:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Docstring mismatch: _sanitize_error claims to remove URLs/stack traces but only truncates

The docstring says "no stack traces or internal URLs", but the implementation only truncates at 500 characters (msg[:500]). If an SDK exception message contains an internal hostname, IP address, or internal URL within the first 500 characters, that information leaks to the MCP client.

Add explicit scrubbing of https?://... URL patterns and File "/path/" stack-trace patterns before the truncation step.

return value.strip()


def _sanitize_path(path: str) -> str:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Docstring mismatch: _sanitize_path claims to reject absolute paths but doesn't

The docstring says "Reject path traversal and absolute paths outside sandbox workspace", but the implementation only checks ".." in p.parts. Absolute paths like /etc/passwd or /proc/1/cmdline pass through. Tool descriptions for sandbox_read_file and sandbox_write_file say "Relative path inside the sandbox workspace".

Either add if p.is_absolute(): raise ValueError(...) or correct the docstring to match.

if name == "sandbox_pause":
try:
await asyncio.to_thread(sb.pause)
return [_text(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Misleading success message: unverified claim about snapshot creation

The message says "A snapshot was created" but the handler only calls sb.pause() without verifying a snapshot was actually produced. If sb.pause() is a pure suspend operation that doesn't create a persistent snapshot, this is misleading.

Either verify the snapshot creation or update the message to accurately describe what sb.pause() does.

if v is not None:
lines.append(f" {k}: {v}")
if not lines:
return [_text(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Leaks API response keys when no known keys are found

If the SDK response has none of the expected keys (sandboxID, templateID, state, cpuCount, memoryMB, startedAt), the full sorted key set is returned to the user (line 475). This leaks internal API structure and undocumented fields.

Return a generic message like "(info returned, no recognised keys)" instead.

@cubesandboxbot

cubesandboxbot Bot commented Jul 7, 2026

Copy link
Copy Markdown

Confirming: review was submitted.

Inline comments on mcp_server.py cover: TOCTOU race with _sandbox.kill() (L524), no-op signal handlers (L555), _sanitize_error URL scrubbing gap (L131), _sanitize_path absolute-path gap (L123), sandbox_pause unverified snapshot claim (L507), sandbox_get_info key leak (L473).

Full review body with all findings was posted above.

loop = asyncio.get_running_loop()
for sig in (signal.SIGTERM, signal.SIGINT):
try:
loop.add_signal_handler(sig, lambda: None)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: No-op signal handlers prevent clean shutdown

The installed handler lambda: None silently consumes SIGTERM and SIGINT instead of initiating shutdown. This means:

  1. SIGINT no longer raises KeyboardInterrupt, so asyncio.run() cannot perform graceful cancellation.
  2. SIGTERM no longer terminates the process (default action is SIG_DFL/terminate).
  3. The atexit-registered _cleanup() never fires, so the sandbox MicroVM remains running as an orphan until its timeout expires (up to 86,400 seconds).

The comment on line 551 claims this forwards signals for "clean shutdown + atexit cleanup", but the no-op achieves the opposite. Replace with a handler that actually stops the event loop, e.g. loop.add_signal_handler(sig, loop.stop), or remove the handlers entirely to rely on default behavior.

return str(p)


def _sanitize_error(exc: BaseException) -> str:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Docstring mismatch: _sanitize_error claims to remove stack traces and internal URLs but only truncates

The docstring says "no stack traces or internal URLs", but the implementation only truncates at 500 characters (msg[:500]). If an exception message (particularly from the SDK's _check_response which may include resp.text from CubeAPI) contains an internal hostname, IP, or internal URL within the first 500 characters, that information is leaked to Claude Code.

Consider adding explicit scrubbing of https?://... URL patterns and File "/path/" stack-trace patterns before truncation.

return value.strip()


def _sanitize_path(path: str) -> str:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Docstring mismatch: _sanitize_path claims to reject absolute paths but doesn't

The docstring says "Reject path traversal and absolute paths outside sandbox workspace", but the implementation only checks ".." in p.parts. Absolute paths like /etc/passwd or /proc/1/cmdline pass through silently. The tool descriptions for sandbox_read_file and sandbox_write_file say "Relative path inside the sandbox workspace", reinforcing the expectation.

Either add if p.is_absolute(): raise ValueError(...) or correct the docstring to match the actual behavior.

if name == "sandbox_pause":
try:
await asyncio.to_thread(sb.pause)
return [_text(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Misleading success message: unverified claim about snapshot creation

The success message says "A snapshot was created and will be accessible via snapshot management" but the handler only calls sb.pause() and does not verify a snapshot was actually created. If sb.pause() is a pure suspend operation that does not produce a persistent snapshot, this message is misleading.

Verify the snapshot creation (e.g., check the return value or call sb.create_snapshot explicitly before pausing), or update the message to accurately describe what sb.pause() does.

if v is not None:
lines.append(f" {k}: {v}")
if not lines:
return [_text(

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Leaks API response keys when no known keys are found

If the SDK returns a response with none of the expected keys (sandboxID, templateID, state, cpuCount, memoryMB, startedAt), the full sorted key set is returned to the user (line 475). This could leak internal API structure or undocumented fields.

Return a generic message like "(info returned, no recognised keys)" instead of exposing the raw key set.

- Fix _sandbox.kill() → sb.kill() (use local var, not stale global)
- _sanitize_path now actually rejects absolute paths
- Fix _sanitize_error docstring to match behaviour
- Add 512KB size limit on sandbox_write_file
- Fix sandbox_pause success message (don't claim snapshot created)

Assisted-by: claude-code:deepseek-v4-pro
- Fix _sandbox.kill() → sb.kill() (use local var for destroy)
- _sanitize_path now actually rejects absolute paths
- Fix _sanitize_error docstring to accurately describe behaviour
- Fix sandbox_pause success message

Assisted-by: claude-code:deepseek-v4-pro
examples/claude-code-sandbox/ → ../sdk/ → wrong
examples/claude-code-sandbox/ → ../../sdk/ → correct

Assisted-by: claude-code:deepseek-v4-pro
# Shutdown — kill sandbox on exit / SIGTERM / SIGINT
# ---------------------------------------------------------------------------
def _cleanup() -> None:
"""Best-effort sandbox kill on process exit."""

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Race condition: _cleanup() reads _sandbox without acquiring _lock

_cleanup() reads _sandbox and calls sb.kill() without holding _lock. The handler registered via atexit may run during shutdown while async tasks are still executing and mutating _sandbox under the lock. This is a data race — a torn read and/or concurrent sb.kill() calls.

Recommendation: Acquire _lock before accessing _sandbox in _cleanup().

min_val=30, max_val=86_400)
except ValueError as exc:
return [_text(f"❌ Invalid argument: {exc}")]

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Missing timeout on Sandbox.create — can permanently consume thread-pool slot

asyncio.to_thread(Sandbox.create, ...) submits the call to a thread-pool. If the CubeAPI endpoint is unreachable or slow, the underlying requests.post() call blocks the thread-pool thread indefinitely (Python stdlib HTTP has no default timeout). Repeated blocked calls would exhaust the thread pool (~5-21 threads), making the server unresponsive to ALL tool calls.

Additionally, the _lock is held during the entire multi-second creation (including the HTTP POST and sandbox boot), blocking other tool requests that queue behind it.

Recommendation: Wrap with asyncio.wait_for(..., timeout=60) and consider releasing the lock during the creation call (re-check _sandbox is None after the call completes).

f" Timeout: {timeout}s"
)]

# ---- guard: sandbox must exist -------------------------------------

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TOCTOU: _sandbox reference captured under lock but used outside it for all tool handlers

sb = _sandbox is read under the lock, then every tool handler (run_code, run_command, read_file, write_file, destroy, etc.) uses sb outside the lock. A concurrent sandbox_destroy call can set _sandbox = None (under the lock) while another handler is mid-way through sb.run_code() or sb.files.read() on a now-destroyed sandbox. The local sb prevents a None dereference, but the underlying API call may fail with a confusing 404 or connection error.

Recommendation: Either re-read _sandbox under the lock before each SDK operation, or use a reader-writer pattern where destroy serializes behind create/run operations.

for the full implementation.

### Custom Template with Extra Tools

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Blocking SDK call inside async function will hang the event loop

This "Minimal MCP Server" snippet shows Sandbox.create(...) called directly inside async def call_tool() without await asyncio.to_thread(...). Since Sandbox.create() is a synchronous blocking call (wraps requests.post), calling it directly in an async function blocks the asyncio event loop, preventing the server from handling any other requests.

The actual mcp_server.py correctly wraps this with await asyncio.to_thread(Sandbox.create, ...). As this is labeled "Minimal" and readers may copy it verbatim, it should use the same non-blocking pattern.

Comment thread examples/claude-code-sandbox/README.md Outdated

- [CubeSandbox Claude Code Integration Guide](../../docs/guide/integrations/claude-code.md) — Full integration guide with best practices
- [CubeSandbox Quick Start](https://cube-sandbox.pages.dev/guide/quickstart)
- [CubeSandbox Python SDK](../sdk/python/README.md)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Broken relative link to Python SDK README

../sdk/python/README.md from inside examples/claude-code-sandbox/ resolves to examples/sdk/python/README.md, which does not exist. The correct relative path should be ../../sdk/python/README.md (go up two levels to repo root, then into sdk/python/).

The same issue exists in README_zh.md at line 254.


## Related

- [CubeSandbox Claude Code Integration](../../../examples/claude-code-sandbox/)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Broken relative link to claude-code-sandbox

../../../examples/claude-code-sandbox/ from examples/pi-agent-sandbox/ resolves one level above the repo root. The correct relative path is ../claude-code-sandbox/ (examples/pi-agent-sandbox/..examples/claude-code-sandbox/).

# Shutdown — kill sandbox on exit / SIGTERM / SIGINT
# ---------------------------------------------------------------------------
def _cleanup() -> None:
"""Best-effort sandbox kill on process exit."""

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Race condition: _cleanup() reads _sandbox without acquiring _lock

_cleanup() at line 65 reads the module-level _sandbox variable and then calls sb.kill() without holding _lock. Since it's registered via atexit, it may run during interpreter shutdown while async tasks are concurrently modifying _sandbox under _lock in call_tool(). This can produce a torn read or a concurrent sb.kill() racing with a destroy call.

Recommendation: Acquire _lock before reading _sandbox in _cleanup(), or mark it as a known limitation with a comment.

…tize URLs

- Add 100KB cap on sandbox_run_code code input
- Add None guard for execution.logs before accessing stdout
- Fix signal handler: use _cleanup() instead of no-op lambda
- _sanitize_error now strips URLs and IP addresses
- Various bot-flagged issue fixes

Assisted-by: claude-code:deepseek-v4-pro
@fslongjin

Copy link
Copy Markdown
Member

Hello, thank you for your contribution! Since we have a large number of participants in this event, to speed up the PR review process, could you please attach some screenshots and logs (both are needed) showing the verification process and results for this PR? Thank you for your participation!

@fslongjin fslongjin left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hello~Can you add more test results/screenshots which can prove this feature can work in a REAL CubeSandbox cluster?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature Request] Integrate CubeSandbox with Pi Agent

3 participants